import numpy as np
from sklearn.model_selection import KFold, StratifiedKFold

x = np.array([
    [1, 2, 3, 4],
    [11, 12, 13, 14],
    [21, 22, 23, 24],
    [31, 32, 33, 34],
    [41, 42, 43, 44],
    [51, 52, 53, 54],
    [61, 62, 63, 64],
    [71, 72, 73, 74],
])

y = np.array([1, 1, 0, 0, 1, 1, 0, 0])

floder = KFold(n_splits=4, shuffle=False)
sFloder = StratifiedKFold(n_splits=4, shuffle=False)

print("KFlold")
for train, test in floder.split(x,y):
    print("train:%s | test:%s",train,test)

print("StratifiedKFold")
for train, test in sFloder.split(x,y):
    print("train:%s | test:%s",train,test)

# 结果：
# KFlold
# train:%s | test:%s [2 3 4 5 6 7] [0 1]
# train:%s | test:%s [0 1 4 5 6 7] [2 3]
# train:%s | test:%s [0 1 2 3 6 7] [4 5]
# train:%s | test:%s [0 1 2 3 4 5] [6 7]
# StratifiedKFold
# train:%s | test:%s [1 3 4 5 6 7] [0 2]
# train:%s | test:%s [0 2 4 5 6 7] [1 3]
# train:%s | test:%s [0 1 2 3 5 7] [4 6]
# train:%s | test:%s [0 1 2 3 4 6] [5 7]